import json
import logging
import re
import threading
import time

import requests

from config import server_address, server_hand_in_address, server_report_progress_address


class Client:
    def __init__(self):
        # 和服务器端通信的会话
        self.server_connection = requests.Session()
        # 和本地webui通信的会话
        self.local_connection = requests.Session()
        # 标识每个客户端的身份
        self.UID = ''
        self.task = {}
        # 传入的格式大概是
        # {
        #   "type": "text2img / img2img",
        #   "prompt": "xxxxxx",
        #   "width": 512,
        #   "height": 512,
        #   "sampler": "Euler a",
        #   "sampling steps": 50,
        # }

        self.task_thread = threading.Thread(target=self.post_drawing_request_to_local, daemon=True)
        self.result = {}
        self.status = ''
        self.percentage = 0.0

    def post_drawing_request_to_local(self):
        self.percentage = 0.0
        s = requests.Session()
        print("绘图线程开始")
        locally_task = {
            "data": [
                self.task['prompt'],
                # self.task['negative_prompt'],
                "",
                "",
                "",
                self.task['sampling steps'],
                self.task['sampler'],
                False,
                False,
                1,
                1,
                7,  # CGF scale 10
                -1,
                -1,
                0,
                0,
                0,
                False,
                self.task['width'],
                self.task['height'],
                False,
                False,
                0.7,  # denoise intensity
                "None",
                False,
                False,
                None,
                "",
                "Seed",
                "",
                "Steps",
                "",
                True,
                False,
                None,
                "",
                ""
            ],
            "fn_index": 12
        }
        try:
            res = s.post("http://localhost:7860/api/predict", json=locally_task)
        except requests.exceptions.ConnectionError as e:
            print(e)
        self.result = res.json()
        self.task = {}
        try:
            self.post_result_to_server()
        except requests.exceptions.ConnectionError as e:
            print(e)
            print('post results failed, possibly the server is down at this moment')

    def post_result_to_server(self):
        meta_data = json.JSONDecoder().decode(self.result['data'][1])
        data = {
            'image': {
                'data_type': self.result['data'][0][0].split(',')[0],
                'raw_data': self.result['data'][0][0].split(',')[1],
            },
            'argument': {
                'prompt': meta_data['prompt'],
                'negative_prompt': meta_data['negative_prompt'],
                'seed': meta_data['seed'],
                'width': meta_data['width'],
                'height': meta_data['height'],
                'sampler': meta_data['sampler'],
                'info_text': meta_data['infotexts']
            }
        }
        self.server_connection.post(server_hand_in_address, json=data)
        self.server_connection.post(server_report_progress_address,
                                    json={"percentage": 100.0}
                                    )

    def send_percentage(self):
        if self.task != {}:
            progress_probe = {
                'fn_index': 2,
                'data': []
            }
            res = self.local_connection.post("http://localhost:7860/api/predict", json=progress_probe)
            res_json = res.json()
            # print(res_json['data'][0])
            regex = re.compile(r"width:(\d{0,}\.\d{0,1})")
            match_group = regex.findall(res_json['data'][0])
            if len(match_group) != 0:
                self.percentage = max(float(match_group[0]), self.percentage)
            self.server_connection.post(server_report_progress_address, json={
                "percentage": self.percentage
            }
                                        )

    def get_server_task(self):
        res = self.server_connection.get(server_address)
        if res.json()['task'] is not None and len(self.task.keys()) == 0:
            self.task = res.json()['task']
            self.task_thread = threading.Thread(target=self.post_drawing_request_to_local, daemon=True)
            self.task_thread.start()

        print(res.json())
        pass


if __name__ == "__main__":
    logging.info("正在启动连接程序...")
    c = Client()
    while True:
        try:
            c.get_server_task()
            c.send_percentage()
        except requests.exceptions.ConnectionError as e:
            print("cannot connect to server, retry in 10 seconds")
            time.sleep(10)
            continue
        time.sleep(1)
